from manim import *

class DrawSquare(Scene):
    def construct(self):
        # 设置边长
        a = 1
        b = 3
        
        # 创建NumberPlane
        plane = NumberPlane()
        self.play(Create(plane))
        
        # 创建大正方形，中心位于原点
        big_square = Square(side_length=a + b, color=GREEN).move_to(ORIGIN)
        
        # 显示大正方形
        self.play(Create(big_square))
        self.wait(2)

        # 获取正方形的四个顶点
        vertices = big_square.get_vertices()
        
        # 存储线段a和线段b的连接点
        connection_points = []

        # 创建并显示标记线段a和线段b
        for i in range(4):
            start = vertices[i]
            end_a = start + (vertices[(i + 1) % 4] - start) * (a / (a + b))
            end_b = vertices[(i + 1) % 4]
            
            # 线段a
            line_a = Line(start, end_a, color=RED)
            self.play(Create(line_a))
            
            # 线段b
            line_b = Line(end_a, end_b, color=YELLOW)
            self.play(Create(line_b))

            # 以brace方式展示a和b
            if i == 0:
                brace_a = Brace(line_a, UP)
                brace_b = Brace(line_b, UP)
            elif i == 1:
                brace_a = Brace(line_a, LEFT)
                brace_b = Brace(line_b, LEFT)
            elif i == 2:
                brace_a = Brace(line_a, DOWN)
                brace_b = Brace(line_b, DOWN)
            else:
                brace_a = Brace(line_a, RIGHT)
                brace_b = Brace(line_b, RIGHT)
            
            label_a = brace_a.get_text("a")
            label_b = brace_b.get_text("b")
            
            # 设置文本颜色
            label_a.set_color(RED)
            label_b.set_color(YELLOW)
            
            self.play(Create(brace_a), Write(label_a))
            self.play(Create(brace_b), Write(label_b))

            # 存储线段a和线段b的连接点
            connection_points.append(end_a)

        self.wait(2)

        # 连接每个边上线段a和线段b的连接点，得到一个正方形，使用颜色绿色
        for i in range(4):
            next_index = (i + 1) % 4
            connection_line = Line(connection_points[i], connection_points[next_index], color=GREEN)
            self.play(Create(connection_line))
        self.wait(2)

        # 上面正方形的边长标记为c
        for i in range(4):
            next_index = (i + 1) % 4
            line_c = Line(connection_points[i], connection_points[next_index])
            brace_c = Brace(line_c, direction=line_c.copy().rotate(PI / 2).get_unit_vector())
            label_c = brace_c.get_text("c").set_color(BLUE)
            self.play(Create(brace_c), Write(label_c))

        self.wait(2)

        # 隐藏NumberPlane
        self.play(FadeOut(plane))

        # 移动所有图形到屏幕左侧
        all_shapes = Group(big_square, *self.mobjects)
        self.play(all_shapes.animate.shift(LEFT * 3))
        self.wait(2)

        # 在屏幕右侧显示如下内容
        # 1. 大正方形的面积 = (a+b)^2
        # 2. 小正方形的面积 = c^2
        # 3. 每个小三角形的面积 = 1/2ab
        # 4. 小正方形的面积 = 大正方形的面积 - 4倍的小三角形面积
        # 5. 推导出 a^2 + b^2 = c^2

        # 简写公式并显示在图形的右侧
        big_square_area = MathTex(
            r"S_{\text{大}} = (a + b)^2 = a^2 + 2ab + b^2",
            tex_template=TexTemplateLibrary.ctex
        ).to_edge(RIGHT).shift(UP * 2)
        self.play(Write(big_square_area))
        self.wait(1)

        small_square_area = MathTex(
            r"S_{\text{小}} = c^2",
            tex_template=TexTemplateLibrary.ctex
        ).next_to(big_square_area, DOWN, buff=0.5)
        self.play(Write(small_square_area))
        self.wait(1)

        triangle_area = MathTex(
            r"S_{\text{三角形}} = \frac{1}{2}ab",
            tex_template=TexTemplateLibrary.ctex
        ).next_to(small_square_area, DOWN, buff=0.5)
        self.play(Write(triangle_area))
        self.wait(1)

        small_square_area_deduction = MathTex(
            r"S_{\text{小}} = S_{\text{大}} - 4 \times S_{\text{三角形}}",
            tex_template=TexTemplateLibrary.ctex
        ).next_to(triangle_area, DOWN, buff=0.5)
        self.play(Write(small_square_area_deduction))
        self.wait(1)

        # 添加动画，把上面的公式代入运算一下，得到最终公式
        final_formula = MathTex(
            r"c^2 = a^2 + 2ab + b^2 - 2ab",
            tex_template=TexTemplateLibrary.ctex
        ).next_to(small_square_area_deduction, DOWN, buff=0.5)
        self.play(Write(final_formula))
        self.wait(1)

        # 推导出 a^2 + b^2 = c^2
        simplified_formula = MathTex(
            r"a^2 + b^2 = c^2",
            tex_template=TexTemplateLibrary.ctex
        ).next_to(small_square_area_deduction, DOWN, buff=0.5)
        self.play(Transform(final_formula, simplified_formula))
        self.wait(2)